"""
零全局变量版 · OpenCV 滑动条实时调参
用法：
    from src.utils.tune_gui import tune_stereo_gui
    best = tune_stereo_gui(left_gray, right_gray)
"""

import cv2
import numpy as np
from dataclasses import dataclass, field
from typing import Callable, Dict
from src.data.structures import BinocularImagePair
from src.algorithms.stereo_matching import SGBMStereoMatcher


@dataclass
class _TunerState:
    """滑条当前值 + 图像缓存"""

    left: np.ndarray
    right: np.ndarray
    win: str
    # 11 个参数默认值
    min_disparity: int = 0
    num_disp_16: int = 14  # 0-16 → 16*(v+1)
    block_size: int = 5  # 奇数
    p1_scale: int = 8  # P1 = scale * channels * block^2
    p2_scale: int = 32
    pre_filter_cap: int = 31
    uniq_ratio: int = 10
    speckle_w: int = 100
    speckle_r: int = 32
    disp12diff: int = 0  # 0 → -1 (禁用)
    mode_idx: int = 2  # 0:SGBM 1:HH 2:3WAY 3:HH4

    # 运行期缓存
    matcher: SGBMStereoMatcher = field(init=False)
    channels: int = field(init=False)

    def __post_init__(self):
        """后置初始化"""
        # 自动推断通道数
        self.channels = 1 if self.left.ndim == 2 else self.left.shape[2]

        # 只创建一次实例
        self.matcher = SGBMStereoMatcher(
            minDisparity=self.min_disparity,
            numDisparities=(self.num_disp_16 + 1) * 16,
            blockSize=self.block_size | 1,
            P1=self.p1_scale * self.channels * self.block_size**2,
            P2=self.p2_scale * self.channels * self.block_size**2,
            preFilterCap=self.pre_filter_cap,
            uniquenessRatio=self.uniq_ratio,
            speckleWindowSize=self.speckle_w,
            speckleRange=self.speckle_r,
            disp12MaxDiff=self.disp12diff - 1,
            mode=[
                cv2.STEREO_SGBM_MODE_SGBM,
                cv2.STEREO_SGBM_MODE_HH,
                cv2.STEREO_SGBM_MODE_SGBM_3WAY,
                cv2.STEREO_SGBM_MODE_HH4,
            ][self.mode_idx],
        )

    # ---------------- 工具方法 ----------------
    def _buildMatcher(self) -> None:
        """根据当前滑条值更新 matcher 参数（不再新建实例）"""
        block = self.block_size | 1
        self.matcher.updateParams(
            min_disparity=self.min_disparity,
            num_disparities=(self.num_disp_16 + 1) * 16,
            block_size=block,
            p1_scale=self.p1_scale,
            p2_scale=self.p2_scale,
            pre_filter_cap=self.pre_filter_cap,
            uniqueness_ratio=self.uniq_ratio,
            speckle_window_size=self.speckle_w,
            speckle_range=self.speckle_r,
            disp12_max_diff=self.disp12diff - 1,
            mode_idx=self.mode_idx,
        )

    def computeAndShow(self) -> None:
        """计算视差并刷新窗口"""
        pair = BinocularImagePair(self.left, self.right)
        disp = self.matcher.computeDisparity(pair)

        disp = np.nan_to_num(disp, nan=0.0)
        disp_vis = (disp - self.min_disparity) / ((self.num_disp_16 + 1) * 16)
        disp_vis = np.clip(disp_vis, 0, 1)
        disp_vis = (disp_vis * 255).astype(np.uint8)
        disp_vis = cv2.applyColorMap(disp_vis, cv2.COLORMAP_JET)

        cv2.imshow(self.win, disp_vis)

    def toDict(self) -> Dict:
        """导出最终参数字典"""
        block = self.block_size | 1
        return {
            "min_disparity": self.min_disparity,
            "num_disparities": (self.num_disp_16 + 1) * 16,
            "block_size": block,
            "p1_scale": self.p1_scale,
            "p2_scale": self.p2_scale,
            "pre_filter_cap": self.pre_filter_cap,
            "uniqueness_ratio": self.uniq_ratio,
            "speckle_window_size": self.speckle_w,
            "speckle_range": self.speckle_r,
            "disp12_max_diff": self.disp12diff - 1,
            "mode_idx": self.mode_idx,
        }


# ------------------------------------------------------------------
# 工厂函数: 把滑条与状态对象绑定
# ------------------------------------------------------------------
def _createTrackbars(state: _TunerState) -> None:
    """创建全部 11 根滑条"""

    def _on(
        name: str,
        max_val: int,
        getter: Callable[[], int],
        setter: Callable[[int], None],
    ):
        """创建滑条回调

        Args:
            name (str): 滑条名称
            max_val (int): 滑条最大值
            getter (Callable[[], int]): 获取当前滑条值的函数
            setter (Callable[[int], None]): 设置当前滑条值的函数
        """

        def callback(pos):
            setter(pos)
            state._buildMatcher()
            state.computeAndShow()

        cv2.createTrackbar(name, state.win, getter(), max_val, callback)

    # 辅助 lambda
    def _simple(name: str, max_val: int, attr: str):
        _on(
            name,
            max_val,
            lambda: getattr(state, attr),
            lambda v: setattr(state, attr, v),
        )

    _simple("minDisparity", 127, "min_disparity")
    _simple("numDisparities(×16)", 16, "num_disp_16")
    _simple("blockSize", 21, "block_size")
    _simple("P1_scale", 50, "p1_scale")
    _simple("P2_scale", 200, "p2_scale")
    _simple("preFilterCap", 63, "pre_filter_cap")
    _simple("uniquenessRatio", 50, "uniq_ratio")
    _simple("speckleWindowSize", 500, "speckle_w")
    _simple("speckleRange", 256, "speckle_r")
    _simple("disp12MaxDiff(+1)", 128, "disp12diff")
    _simple("mode(0-SGBM 1-HH 2-3WAY 3-HH4)", 3, "mode_idx")


# ------------------------------------------------------------------
# 对外入口
# ------------------------------------------------------------------
def tuneSgbmGui(
    rectified_left: np.ndarray, rectified_right: np.ndarray, default: dict | None = None
) -> Dict:
    """弹窗调参, 返回最优参数字典"""
    if rectified_left.shape != rectified_right.shape:
        raise ValueError("左右图像尺寸不匹配")
    if rectified_left.ndim not in (2, 3):
        raise ValueError("仅支持灰度图或彩色图")

    win = "SGBM Tuner"
    cv2.namedWindow(win, cv2.WINDOW_KEEPRATIO)
    cv2.resizeWindow(win, 900, 900)

    state = _TunerState(
        left=rectified_left, right=rectified_right, win=win, **(default or {})
    )
    _createTrackbars(state)

    state.computeAndShow()  # 初始图
    print("【使用帮助】")
    print("  - 拖动任意滑条 → 实时刷新视差彩图")
    print("  - 按 ESC 退出并返回当前参数字典")

    while True:
        if cv2.waitKey(1) & 0xFF == 27:  # ESC
            break
    cv2.destroyAllWindows()
    return state.toDict()
